import os
import utils
from paddle.io import Dataset

class pokemonDataset(Dataset):
    
    def __init__(self,dataset_path = ""):
        self.dataset_path = dataset_path
        self.img_path = os.listdir(dataset_path)
        
    def __getitem__(self,index):
        img = utils.read_bgr(os.path.join(self.dataset_path,self.img_path[index]))
        img = utils.channel_last_to_first(img)
        x = y = img/255.
        return x,y
    
    def  __len__(self):
        return len(self.img_path)